#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
最小化ASR语音识别服务

基于火山引擎离线ASR服务的最小可复用单元
提供核心的语音识别功能，去除复杂的文本重对齐逻辑
"""

import json
import time
import uuid
import hashlib
import requests
from pathlib import Path


class MinimalASR:
    """
    最小化ASR服务 - 保留核心功能的可复用单元
    
    核心功能：
    1. 语音识别API调用
    2. 基本缓存机制
    3. 结果解析
    """
    
    def __init__(self, app_id=None, access_token=None, cache_dir="./asr_cache"):
        """
        初始化ASR服务
        
        Args:
            app_id: 火山引擎应用ID
            access_token: 访问令牌
            cache_dir: 缓存目录路径
        """
        # 火山ASR API配置 - 支持外部传入
        self.app_id = app_id or "3967329463"
        self.access_token = access_token or "GibkwaUNpMNKNMfwAn0r1uL0CHm6gDDs"
        
        # API端点
        self.submit_url = "https://openspeech.bytedance.com/api/v3/auc/bigmodel/submit"
        self.query_url = "https://openspeech.bytedance.com/api/v3/auc/bigmodel/query"
        
        # 状态码
        self.STATUS_CODE_OK = 20000000
        self.STATUS_CODE_PROCESSING = 20000001
        self.STATUS_CODE_IN_TASK_QUEUE = 20000002
        
        # 缓存配置
        self.cache_dir = Path(cache_dir)
        self.cache_dir.mkdir(exist_ok=True)
        self.enable_cache = True
        
        print(f"ASR服务初始化完成，缓存目录: {self.cache_dir}")
    
    def _get_cache_key(self, audio_url, identifier=None):
        """生成缓存键"""
        if identifier:
            base_key = str(identifier)
        else:
            # 使用URL的hash作为缓存键
            base_key = hashlib.md5(audio_url.encode()).hexdigest()[:12]
        return f"asr_{base_key}"
    
    def _get_cache_file(self, cache_key):
        """获取缓存文件路径"""
        return self.cache_dir / f"{cache_key}.json"
    
    def _load_cache(self, audio_url, identifier=None):
        """加载缓存结果"""
        if not self.enable_cache:
            return None
            
        cache_key = self._get_cache_key(audio_url, identifier)
        cache_file = self._get_cache_file(cache_key)
        
        if cache_file.exists():
            try:
                with open(cache_file, 'r', encoding='utf-8') as f:
                    cached_data = json.load(f)
                print(f"从缓存加载ASR结果: {cache_key}")
                return cached_data.get('segments', [])
            except Exception as e:
                print(f"加载缓存失败: {e}")
        return None
    
    def _save_cache(self, audio_url, segments, identifier=None):
        """保存结果到缓存"""
        if not self.enable_cache:
            return False
            
        cache_key = self._get_cache_key(audio_url, identifier)
        cache_file = self._get_cache_file(cache_key)
        
        cache_data = {
            'audio_url': audio_url,
            'identifier': identifier,
            'cached_time': time.strftime('%Y-%m-%d %H:%M:%S'),
            'segments': segments,
            'total_segments': len(segments)
        }
        
        try:
            with open(cache_file, 'w', encoding='utf-8') as f:
                json.dump(cache_data, f, ensure_ascii=False, indent=2)
            print(f"结果已缓存: {cache_key}")
            return True
        except Exception as e:
            print(f"保存缓存失败: {e}")
            return False
    
    def _submit_asr_request(self, audio_url, request_id):
        """提交ASR识别请求"""
        request_body = {
            "user": {"uid": "minimal_asr_user"},
            "audio": {
                "url": audio_url,
                "format": "wav",
                "codec": "raw", 
                "rate": 16000,
                "bits": 16,
                "channel": 1
            },
            "request": {"model_name": "bigmodel"}
        }
        
        headers = {
            "Content-Type": "application/json",
            "X-Api-App-Key": self.app_id,
            "X-Api-Access-Key": self.access_token,
            "X-Api-Resource-Id": "volc.bigasr.auc",
            "X-Api-Request-Id": request_id,
            "X-Api-Sequence": "-1"
        }
        
        try:
            response = requests.post(self.submit_url, headers=headers, json=request_body, timeout=60)
            status_code = int(response.headers.get("X-Api-Status-Code", "-1"))
            
            if status_code == self.STATUS_CODE_OK:
                return response.headers.get("X-Tt-Logid", "")
            else:
                print(f"ASR提交失败: {response.headers.get('X-Api-Message', '')}")
                return None
        except Exception as e:
            print(f"ASR提交异常: {e}")
            return None
    
    def _query_asr_result(self, request_id, max_attempts=60, poll_interval=5):
        """查询ASR识别结果（轮询）"""
        headers = {
            "Content-Type": "application/json",
            "X-Api-App-Key": self.app_id,
            "X-Api-Access-Key": self.access_token,
            "X-Api-Resource-Id": "volc.bigasr.auc",
            "X-Api-Request-Id": request_id
        }
        
        for attempt in range(max_attempts):
            if attempt > 0:
                time.sleep(poll_interval)
            
            try:
                response = requests.post(self.query_url, headers=headers, json={}, timeout=30)
                status_code = int(response.headers.get("X-Api-Status-Code", "-1"))
                
                if status_code == self.STATUS_CODE_OK:
                    return response.text
                elif status_code in [self.STATUS_CODE_PROCESSING, self.STATUS_CODE_IN_TASK_QUEUE]:
                    if attempt % 10 == 0:  # 每50秒打印一次状态
                        print(f"⏳ ASR处理中... ({attempt * poll_interval}s)")
                    continue
                else:
                    print(f"ASR查询失败: {response.headers.get('X-Api-Message', '')}")
                    return None
            except Exception as e:
                print(f"查询异常: {e}")
                if attempt == max_attempts - 1:
                    return None
        
        print(f"ASR查询超时")
        return None
    
    def _parse_asr_result(self, asr_result_json):
        """
        解析火山ASR识别结果
        
        Args:
            asr_result_json: ASR识别原始结果JSON字符串
            
        Returns:
            list: 语音段列表，每个段包含 start, end, text, words
        """
        try:
            result_data = json.loads(asr_result_json)
            
            # 尝试主要路径: result.utterances
            utterances = None
            if 'result' in result_data and 'utterances' in result_data['result']:
                utterances = result_data['result']['utterances']
            elif 'data' in result_data and 'utterances' in result_data['data']:
                utterances = result_data['data']['utterances']
            
            if not utterances:
                print(f"未找到utterances数据")
                return []
            
            segments = []
            for utterance in utterances:
                if not all(key in utterance for key in ['text', 'start_time', 'end_time']):
                    continue
                
                # 时间转换：毫秒 -> 秒
                start_time = utterance['start_time'] / 1000.0
                end_time = utterance['end_time'] / 1000.0
                text = utterance['text'].strip()
                
                if not text or end_time <= start_time:
                    continue
                
                # 解析words信息
                words = []
                if 'words' in utterance and isinstance(utterance['words'], list):
                    for word_data in utterance['words']:
                        if isinstance(word_data, dict) and 'text' in word_data:
                            word = {
                                'text': word_data.get('text', '').strip(),
                                'start_time': word_data.get('start_time', 0) / 1000.0,
                                'end_time': word_data.get('end_time', 0) / 1000.0,
                                'confidence': word_data.get('confidence', 0.0)
                            }
                            if word['text'] and word['end_time'] > word['start_time']:
                                words.append(word)
                
                segment = {
                    "start": start_time,
                    "end": end_time,
                    "text": text,
                    "words": words
                }
                segments.append(segment)
            
            print(f"解析完成: {len(segments)} 个语音段")
            return segments
            
        except Exception as e:
            print(f"ASR结果解析失败: {e}")
            return []
    
    def recognize(self, audio_url, identifier=None, use_cache=True):
        """
        语音识别主接口
        
        Args:
            audio_url: 音频文件的远程URL
            identifier: 可选的标识符，用于缓存键生成
            use_cache: 是否使用缓存
            
        Returns:
            list: 语音段列表，每个段包含 start, end, text, words
        """
        print(f"开始ASR识别:")
        print(f"    音频URL: {audio_url}")
        print(f"    标识符: {identifier}")
        
        # 检查缓存
        if use_cache and self.enable_cache:
            cached_result = self._load_cache(audio_url, identifier)
            if cached_result:
                return cached_result
        
        # 生成请求ID并提交请求
        request_id = str(uuid.uuid4())
        print(f"    请求ID: {request_id}")
        
        # 提交ASR请求
        logid = self._submit_asr_request(audio_url, request_id)
        if not logid:
            print(f"ASR请求提交失败")
            return []
        
        # 查询结果
        asr_result_json = self._query_asr_result(request_id)
        if not asr_result_json:
            print(f"ASR结果查询失败")
            return []
        
        # 解析结果
        segments = self._parse_asr_result(asr_result_json)
        
        # 保存到缓存
        if segments and use_cache:
            self._save_cache(audio_url, segments, identifier)
        
        print(f"ASR识别完成: {len(segments)} 个语音段")
        return segments
    
    def get_full_text(self, segments):
        """
        从语音段中提取完整文本
        
        Args:
            segments: 语音段列表
            
        Returns:
            str: 完整的识别文本
        """
        return " ".join(segment['text'] for segment in segments if segment.get('text'))
    
    def get_words_with_timestamps(self, segments):
        """
        从语音段中提取所有词及其时间戳
        
        Args:
            segments: 语音段列表
            
        Returns:
            list: 词列表，每个词包含 text, start_time, end_time, confidence
        """
        all_words = []
        for segment in segments:
            if 'words' in segment:
                all_words.extend(segment['words'])
        return all_words
    
    def disable_cache(self):
        """禁用缓存"""
        self.enable_cache = False
        print("缓存已禁用")
    
    def enable_cache_func(self):
        """启用缓存"""
        self.enable_cache = True
        print("缓存已启用")
    
    def clear_cache(self):
        """清空缓存目录"""
        try:
            for cache_file in self.cache_dir.glob("asr_*.json"):
                cache_file.unlink()
            print("缓存已清空")
        except Exception as e:
            print(f"清空缓存失败: {e}")


# 使用示例
if __name__ == "__main__":
    # 初始化ASR服务
    asr = MinimalASR()
    
    # 示例音频URL（需要替换为实际的音频URL）
    audio_url = "/Users/songyuliang/Desktop/北京语言大学\ 3.m4a"
    
    # 进行语音识别
    segments = asr.recognize(audio_url, identifier="test_audio")
    
    # 输出结果
    if segments:
        print(f"\n识别结果:")
        for i, segment in enumerate(segments):
            print(f"  {i+1}. [{segment['start']:.2f}s - {segment['end']:.2f}s] {segment['text']}")
        
        # 获取完整文本
        full_text = asr.get_full_text(segments)
        print(f"\n完整文本: {full_text}")
        
        # 获取词级时间戳
        words = asr.get_words_with_timestamps(segments)
        print(f"\n词数统计: {len(words)} 个词")
    else:
        print("识别失败")